import subprocess
import itertools
import os
from datetime import datetime

base_command_parts = [
    "PYTHONUNBUFFERED=1",
    "python3", "-m", "verl.trainer.main_ppo",
    "data.train_files=/datasets/convert_gsm8k/train.parquet",
    "data.val_files=/datasets/convert_gsm8k/test.parquet",
    "data.train_batch_size=128",
    "data.max_prompt_length=1024",
    "data.max_response_length=2048",
    "actor_rollout_ref.rollout.name=sglang",
    "actor_rollout_ref.model.path=/data/zhujunjie/Llama-3.1-8B-Instruct",
    "actor_rollout_ref.actor.optim.lr=1e-6",
    "actor_rollout_ref.actor.ppo_mini_batch_size=128",
    "actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16",
    "actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16",
    "actor_rollout_ref.model.enable_gradient_checkpointing=True",
    "actor_rollout_ref.actor.fsdp_config.param_offload=True",
    "actor_rollout_ref.actor.fsdp_config.optimizer_offload=True",
    "actor_rollout_ref.rollout.tensor_model_parallel_size=2",
    "actor_rollout_ref.rollout.gpu_memory_utilization=0.3", # 调整这个值如果显存不足
    "actor_rollout_ref.rollout.dtype=float16",
    "+actor_rollout_ref.rollout.speculative_algorithm=EAGLE3",
    "+actor_rollout_ref.rollout.speculative_draft_model_path=/data/zhujunjie/sglang-EAGLE3-Llama-3.1-Instruct-8B",
    "+actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16",
    "critic.optim.lr=1e-5",
    "critic.model.path=/data/zhujunjie/Llama-3.1-8B-Instruct", # Critic 模型路径
    "critic.ppo_micro_batch_size_per_gpu=16",
    "critic.model.fsdp_config.param_offload=True",
    "critic.model.fsdp_config.optimizer_offload=True",
    "algorithm.kl_ctrl.kl_coef=0.001",
    "trainer.logger=['console']",
    "trainer.val_before_train=False",
    "trainer.default_hdfs_dir=null",
    "trainer.n_gpus_per_node=8",
    "trainer.nnodes=1",
    "trainer.save_freq=-1",
    "trainer.test_freq=10", 
    "trainer.total_epochs=1" 
]

# 投机解码参数的选项列表
speculative_num_steps_options = [3, 5]
speculative_eagle_topk_options = [4, 8]
speculative_num_draft_tokens_options = [8, 32]

# 日志文件输出目录
log_output_dir = "verl_speculative_logs"
os.makedirs(log_output_dir, exist_ok=True)

# 设置 SGLang 环境变量 (如果不在 base_command_parts 中)
env_vars = {"SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK": "True"}
# 如果 PYTHONUNBUFFERED=1 需要在命令前，可以这样处理
# 或者你也可以在 base_command_parts 第一个元素就是 "PYTHONUNBUFFERED=1"
# 然后在 subprocess.run 中通过 shell=True 执行，或者分别处理环境变量

# --- 网格搜索并执行 ---
experiment_count = 0
total_experiments = len(speculative_num_steps_options) * \
                  len(speculative_eagle_topk_options) * \
                  len(speculative_num_draft_tokens_options)

print(f"Starting grid search for {total_experiments} experiments...\n")

for num_steps in speculative_num_steps_options:
    for eagle_topk in speculative_eagle_topk_options:
        for num_draft_tokens in speculative_num_draft_tokens_options:
            experiment_count += 1
            print(f"--- Running Experiment {experiment_count}/{total_experiments} ---")
            print(f"Config: num_steps={num_steps}, eagle_topk={eagle_topk}, num_draft_tokens={num_draft_tokens}")

            # 构建当前实验的特定参数
            current_speculative_params = [
                f"+actor_rollout_ref.rollout.speculative_num_steps={num_steps}",
                f"+actor_rollout_ref.rollout.speculative_eagle_topk={eagle_topk}",
                f"+actor_rollout_ref.rollout.speculative_num_draft_tokens={num_draft_tokens}"
            ]

            log_filename = os.path.join(
                log_output_dir,
                f"verl_s{num_steps}_k{eagle_topk}_d{num_draft_tokens}.log"
            )

            # 组合完整命令
            # 将PYTHONUNBUFFERED=1作为命令的一部分，使用空格连接
            full_command_list = base_command_parts[:1] + [" ".join(base_command_parts[1:3])] + base_command_parts[3:] + current_speculative_params
            full_command_str = " ".join(full_command_list)
            
            command_to_run = f"{full_command_str} 2>&1 | tee {log_filename}"

            print(f"Executing: {command_to_run}\n")

            try:
                # 获取当前环境变量，并更新
                current_env = os.environ.copy()
                current_env.update(env_vars)

                process = subprocess.Popen(command_to_run, shell=True, env=current_env)
                process.wait() # 等待当前实验完成

                if process.returncode == 0:
                    print(f"Experiment {experiment_count} completed successfully. Log: {log_filename}")
                else:
                    print(f"Experiment {experiment_count} FAILED with return code {process.returncode}. Log: {log_filename}")
            except Exception as e:
                print(f"Experiment {experiment_count} CRASHED with exception: {e}. Log: {log_filename}")
            
            print("-" * 50 + "\n")

print("All experiments finished.")